import os
import logging
from flask import redirect, flash, render_template, url_for
# 从werkzeug.utils导入secure_filename函数，用于安全处理文件名
from werkzeug.utils import secure_filename
# 从utils模块导入allowed_doc_file函数，用于判断文件类型是否合法
from utils import allowed_doc_file, extract_doc_content
from embedding_utils import get_doc_embedding
# 从milvus_utils模块导入插入向量的函数
from milvus_utils import insert_doc_vectors, search_doc_vectors

logger = logging.getLogger(__name__)

UPLOAD_FOLDER = os.getenv('UPLOAD_FOLDER', 'uploads')
if not os.path.exists(UPLOAD_FOLDER):
    os.makedirs(UPLOAD_FOLDER)

def handle_doc_upload(request):
    # 如果请求方法为POST
    if request.method == "POST":
        logger.info("文档上传POST请求")
        if "file" not in request.files:
            logger.warning("文件名为空")
            flash("文件名为空", "error")
            return redirect(request.url)
        file = request.files["file"]
        # 检查文件名是否为空
        if file.filename == "":
            # 记录文件名为空的日志
            logger.warning("未选择文件")
            # 弹出提示信息
            flash("未选择文件", "error")
            # 重定向回当前页面
            return redirect(request.url)

        # 如果文件存在且文件类型合法
        if file and allowed_doc_file(file.filename):
            # 文件名安全处理
            filename = secure_filename(file.filename)
            # 拼接文件路径
            file_path = os.path.join(UPLOAD_FOLDER, filename)
            # 保存文件到指定路径
            file.save(file_path)
            logger.info(f"文件已保存到：{file_path}")
            # 提取文档内容
            doc_content = extract_doc_content(file_path)
            logger.info(f"文件内容提取成功，内容长度：{len(doc_content)}")
            # 计算文档向量
            doc_embedding = get_doc_embedding(doc_content)
            logger.info(f"文档向量计算成功，向量长度{len(doc_embedding)}")
            #插入文档向量和原文档内容到Milvus数据库
            insert_doc_vectors([doc_embedding], [doc_content])
            flash("文件上传成功", "success")
            # 重定向到上传文档页面
            return redirect(url_for("upload_doc"))

    # 渲染文档页面
    return render_template("upload_doc.html")

def handle_doc_search(request):
    results = []
    query = ""
    if request.method == "POST":
        logger.info("文档搜索POST请求")
        # 获取表单中的检索文本
        query = request.form.get("query")
        logger.info(f"检索文本{query}")
        if query.strip():
            try:
                logger.info("开始文本向量化检索")
                # 获取检索文本的向量表示
                query_bedding = get_doc_embedding(query)
                logger.info("文本向量化成功，开始Milvus检索")
                # 在Milvus中检索最相似的文档
                milvus_results = search_doc_vectors(query_bedding, top_k=5)
                logger.info(f"Milvus检索完成，命中组数：{len(milvus_results)}")
                # 遍历检索结果
                for hits in milvus_results:
                    for hit in hits:
                        # 将每个命中结果添加到结果列表
                        results.append({
                            "doc": hit.entity.get("doc", ""),
                            "score": hit.distance,
                        })
            except Exception as e:
                logger.error(f"检索失败：{e}")
                # 闪现检索失败的提示信息
                flash(f"检索失败：{e}")
    # 记录检索结果
    logger.info(f"检索结果{results}")
    return render_template("search_doc.html", results=results, query=query)